#include "kernel_operator.h"
#include "kernel_utils.h"
#include "lib/matmul_intf.h"
#include "utils.h"
using namespace AscendC;
using namespace matmul;

template <typename T> class KernelMatmulSub {
public:
  __aicore__ inline KernelMatmulSub() {}
  __aicore__ inline void
  Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR x3, GM_ADDR y, GM_ADDR workspace,
       const MatMulSubTilingData *__restrict tiling, TPipe *pipe);
  __aicore__ inline void Process();
  __aicore__ inline void DoMatmul();
  __aicore__ inline void DoSub();

  MatmulImpl<MatmulType<TPosition::GM, CubeFormat::ND, T, false>,
             MatmulType<TPosition::GM, CubeFormat::ND, T, false>,
             MatmulType<TPosition::GM, CubeFormat::ND, T>,
             MatmulType<TPosition::GM, CubeFormat::ND, T>, CFG_MDL>
      mmObj;

  struct MNConfig {
    uint32_t m;
    uint32_t n;
    uint32_t k;
    uint32_t baseM;
    uint32_t baseN;
    uint32_t blockDimM;
    uint32_t blockDimN;
    uint32_t mIdx;
    uint32_t nIdx;
    uint32_t singleM;
    uint32_t singleN;
  } mnConfig;

  int32_t iterateNum;
  uint32_t coreIdx;
  uint32_t subBlockIdx;
  uint32_t blockNum;
  const MatMulSubTilingData *__restrict tilingData;
  GlobalTensor<T> x1Global;
  GlobalTensor<T> x2Global;
  GlobalTensor<T> x3Global;
  GlobalTensor<T> yGlobal;
  GlobalTensor<T> mmOutGlobal;

  // AIV
  uint32_t tileSizeM;
  LocalTensor<T> mmLocal;
  LocalTensor<T> x3Local;
  LocalTensor<T> outLocal;
  TQue<QuePosition::VECIN, 1> mmQueue;
  TQue<QuePosition::VECIN, 1> x3Queue;
  TQue<QuePosition::VECOUT, 1> outQueue;
};

template <typename T>
__aicore__ inline void KernelMatmulSub<T>::Init(
    GM_ADDR x1, GM_ADDR x2, GM_ADDR x3, GM_ADDR y, GM_ADDR workspace,
    const MatMulSubTilingData *__restrict tiling, TPipe *pipe) {

  iterateNum = 0;
  coreIdx = GetBlockIdx() / GetTaskRation();
  subBlockIdx = GetSubBlockIdx();
  tilingData = tiling;
  x1Global.SetGlobalBuffer((__gm__ T *)x1);
  x2Global.SetGlobalBuffer((__gm__ T *)x2);
  x3Global.SetGlobalBuffer((__gm__ T *)x3);
  yGlobal.SetGlobalBuffer((__gm__ T *)y);
  mmOutGlobal.SetGlobalBuffer((__gm__ T *)workspace);

  mnConfig.m = tilingData->M;
  mnConfig.n = tilingData->N;
  mnConfig.k = tilingData->K;
  mnConfig.singleM = tilingData->singleM;
  mnConfig.singleN = tilingData->singleN;
  mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
  mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
  blockNum = mnConfig.blockDimM * mnConfig.blockDimN;

  if ASCEND_IS_AIV {
    tileSizeM = (tilingData->dimX3 == 1) ? 1 : (16 / sizeof(T));
    pipe->InitBuffer(mmQueue, 2, tileSizeM * tilingData->singleN * sizeof(T));
    pipe->InitBuffer(x3Queue, 2, tileSizeM * tilingData->singleN * sizeof(T));
    pipe->InitBuffer(outQueue, 2, tileSizeM * tilingData->singleN * sizeof(T));
  }
}

template <typename T> __aicore__ inline void KernelMatmulSub<T>::Process() {
  uint32_t KN = tilingData->K * tilingData->N;
  uint32_t blockIdx = coreIdx;
  while (blockIdx < blockNum) {
    mnConfig.mIdx = blockIdx / mnConfig.blockDimN;
    mnConfig.nIdx = blockIdx % mnConfig.blockDimN;
    DoMatmul();
    DoSub();
    blockIdx += tilingData->usedCoreNum;
    iterateNum += 1;
  }
  if ASCEND_IS_AIC {
    if (likely(iterateNum > 0)) {
      for (int32_t idx = 0; idx <= (iterateNum - 1) % PIPELINE_NUM; ++idx) {
        WaitAiv();
      }
    }
  }
}

template <typename T> __aicore__ inline void KernelMatmulSub<T>::DoMatmul() {
  if ASCEND_IS_AIV {
    return;
  }
  uint32_t nIdx = mnConfig.nIdx * mnConfig.singleN;
  uint32_t mIdx = mnConfig.mIdx * mnConfig.singleM;
  uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1
                            ? mnConfig.singleN
                            : mnConfig.n - nIdx;
  uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1
                            ? mnConfig.singleM
                            : mnConfig.m - mIdx;
  mmObj.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
  mmObj.SetSingleShape(curSingleM, curSingleN, mnConfig.k);
  mmObj.SetTensorA(x1Global[mIdx * mnConfig.k]);
  mmObj.SetTensorB(x2Global[nIdx]);
  if (iterateNum >= PIPELINE_NUM) {
    WaitAiv();
  }
  mmObj.IterateAll(mmOutGlobal[mIdx * mnConfig.n + nIdx]);
  NotifyAiv();
}

template <typename T> __aicore__ inline void KernelMatmulSub<T>::DoSub() {
  if ASCEND_IS_AIC {
    return;
  }
  uint32_t nIdx = mnConfig.nIdx * mnConfig.singleN;
  uint32_t mIdx = mnConfig.mIdx * mnConfig.singleM;
  uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1
                            ? mnConfig.singleN
                            : mnConfig.n - nIdx;
  uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1
                            ? mnConfig.singleM
                            : mnConfig.m - mIdx;
  uint32_t alignCurSingleN = AlignUp(
      curSingleN, static_cast<uint32_t>(UB_BLOCK_DOUBLE_UNIT_SIZE / sizeof(T)));
  uint32_t curSingleMOffset = 0;
  if (subBlockIdx == 0) {
    curSingleM = curSingleM / 2;
  } else {
    curSingleMOffset = curSingleM / 2;
    curSingleM = curSingleM - curSingleMOffset;
  }
  uint32_t mIdxGlobal = mIdx + curSingleMOffset;
  uint32_t outOffset = mIdxGlobal * mnConfig.n + nIdx;

  WaitAic();
  NotifyAic();

  if(tilingData->dimX3 == 1) {
    x3Local = x3Queue.AllocTensor<T>();
    DataCopyPad2D(x3Local, x3Global[nIdx], 1,
                  curSingleN, mnConfig.n);
    x3Queue.EnQue(x3Local);
    x3Local = x3Queue.DeQue<T>();
  }

  for (int offsetM = 0; offsetM < curSingleM; offsetM += tileSizeM) {
    uint32_t sizeM = (offsetM + tileSizeM >= curSingleM)
                         ? (curSingleM - offsetM)
                         : tileSizeM;
    mmLocal = mmQueue.AllocTensor<T>();
    DataCopyPad2D(mmLocal, mmOutGlobal[outOffset + offsetM * mnConfig.n], sizeM,
                  curSingleN, mnConfig.n);
    mmQueue.EnQue(mmLocal);
    mmLocal = mmQueue.DeQue<T>();
    if(tilingData->dimX3 != 1) {
      x3Local = x3Queue.AllocTensor<T>();
      DataCopyPad2D(x3Local, x3Global[outOffset + offsetM * mnConfig.n], sizeM,
                  curSingleN, mnConfig.n);
      x3Queue.EnQue(x3Local);
      x3Local = x3Queue.DeQue<T>();
    }
    outLocal = outQueue.AllocTensor<T>();
    outLocal = mmLocal - x3Local;
    outQueue.EnQue(outLocal);
    outLocal = outQueue.DeQue<T>();
    DataCopyPad2D(yGlobal[outOffset + offsetM * mnConfig.n], outLocal, sizeM,
                  curSingleN, alignCurSingleN, mnConfig.n);
    mmQueue.FreeTensor(mmLocal);
    outQueue.FreeTensor(outLocal);
    if(tilingData->dimX3 != 1) {
      x3Queue.FreeTensor(x3Local);
    }
  }
  if(tilingData->dimX3 == 1) {
    x3Queue.FreeTensor(x3Local);
  }
}